package com.augrain.microframework.limit;

import com.augrain.microframework.limit.aspect.RateLimitAspectHandler;
import com.augrain.microframework.limit.config.RateLimitConfig;
import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.context.annotation.Import;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.data.redis.core.script.RedisScript;
import org.springframework.data.redis.serializer.StringRedisSerializer;

/**
 * RateLimit自动装配
 *
 * @author biaoy
 * @since 2020/12/22
 */
@Configuration
@ConditionalOnProperty(prefix = RateLimitConfig.PREFIX, name = "enable", havingValue = "true", matchIfMissing = true)
@Import({RateLimitAspectHandler.class, RateLimitConfig.class})
public class RateLimitAutoConfiguration {

    @Bean
    @ConditionalOnMissingBean
    public LimitKeyGenerator limitKeyGenerator() {
        return new LimitKeyGenerator();
    }

    @Bean
    public RedisScript<String> redisScript() {
        String script = "redis.replicate_commands()\n" +
                "\n" +
                "local key = KEYS[1] -- 令牌桶标识\n" +
                "local quantity = 1 -- 需要的令牌数量，默认为1\n" +
                "local timestamp = tonumber(redis.call('time')[1]) -- 当前时间戳\n" +
                "local result = true\n" +
                "local expire = -1 -- 过期时间按照多规则中最大的时间来处理\n" +
                "\n" +
                "-- 第一次请求时创建令牌桶\n" +
                "if (redis.call('EXISTS', key) == 0) then\n" +
                "    for i = 1, #ARGV, 2 do\n" +
                "        local quota = tonumber(ARGV[i])\n" +
                "        local period = tonumber(ARGV[i + 1])\n" +
                "        expire = period\n" +
                "\n" +
                "        local remainKey = 'r' .. i\n" +
                "        local timeKey = 't' .. i\n" +
                "        redis.call('HMSET', key, remainKey, quota - quantity, timeKey, timestamp + period)\n" +
                "    end\n" +
                "else\n" +
                "    for i = 1, #ARGV, 2 do\n" +
                "        local quota = tonumber(ARGV[i])  -- 时间窗口内的限额\n" +
                "        local period = tonumber(ARGV[i + 1]) -- 时间窗口大小（秒）\n" +
                "        local remainKey = 'r' .. i\n" +
                "        local timeKey = 't' .. i\n" +
                "        expire = period\n" +
                "\n" +
                "        -- 计算从上次生成到现在这段时间应该生成的令牌数\n" +
                "        local remain = tonumber(redis.call('HGET', key, remainKey))\n" +
                "        local last_reset = tonumber(redis.call('HGET', key, timeKey))\n" +
                "        local delta_quota = math.floor(((timestamp - last_reset) / period) * quota)\n" +
                "        if (delta_quota > 0) then\n" +
                "            remain = remain + delta_quota\n" +
                "            if (remain > quota) then\n" +
                "                remain = quota\n" +
                "            end\n" +
                "            redis.call('HMSET', key, remainKey, remain, timeKey, timestamp)\n" +
                "        end\n" +
                "        -- 如果剩余令牌数量大于请求数量，则该规则下获取令牌成功\n" +
                "        if (remain < quantity) then\n" +
                "            result = false\n" +
                "        end\n" +
                "    end\n" +
                "\n" +
                "    -- 如果所有规则都成功了\n" +
                "    if (result == true) then\n" +
                "        -- 则更新剩余token\n" +
                "        for i = 1, #ARGV, 2 do\n" +
                "            local remainKey = 'r' .. i\n" +
                "            redis.call('HINCRBY', key, remainKey, -quantity)\n" +
                "        end\n" +
                "    end\n" +
                "end\n" +
                "redis.call('EXPIRE', key, expire)\n" +
                "return tostring(result)";

        DefaultRedisScript<String> redisScript = new DefaultRedisScript<>(script);
        redisScript.setResultType(String.class);
        return redisScript;
    }

    @Bean
    public StringRedisSerializer stringRedisSerializer() {
        return new StringRedisSerializer();
    }
}
